from math import floor
from types import SimpleNamespace


def ax_hyparams(hyp_params, search_hyp_params_list):
  len_dict = {
      'l': hyp_params.l_len,
      'a': hyp_params.a_len,
      'v': hyp_params.v_len
  }
  bound_dict = {}
  for name_i, len_i in len_dict.items():
    for name_j, len_j in len_dict.items():
      # generate upper limit, consider the order first.
      ks_base = floor(len_j / len_i) // 2 * 2
      upper_bound = min(len_j - 1, 49)
      search_bound = (upper_bound - ks_base) // 2
      bound_dict[name_i + name_j] = search_bound

  ax_hyparams_config_base = {
      "kernel_size_av": {
          "name": "kernel_size_av",
          "type": "range",
          "bounds": [1, bound_dict['av']],
          "value_type": "int",
          "log_scale": False,
      },
      "kernel_size_al": {
          "name": "kernel_size_al",
          "type": "range",
          "bounds": [1, bound_dict['al']],
          "value_type": "int",
          "log_scale": False,
      },
      "kernel_size_la": {
          "name": "kernel_size_la",
          "type": "range",
          "bounds": [1, bound_dict['la']],
          "value_type": "int",
          "log_scale": False,
      },
      "kernel_size_lv": {
          "name": "kernel_size_lv",
          "type": "range",
          "bounds": [1, bound_dict['lv']],
          "value_type": "int",
          "log_scale": False,
      },
      "kernel_size_va": {
          "name": "kernel_size_va",
          "type": "range",
          "bounds": [1, bound_dict['va']],
          "value_type": "int",
          "log_scale": False,
      },
      "kernel_size_vl": {
          "name": "kernel_size_vl",
          "type": "range",
          "bounds": [1, bound_dict['vl']],
          "value_type": "int",
          "log_scale": False,
      },
      "kernel_size": {
          "name": "kernel_size",
          "type": "range",
          "bounds": [1, min(bound_dict.values())],
          "value_type": "int",
          "log_scale": False,
      },
      "kernel_size_v": {
          "name": "kernel_size_v",
          "type": "range",
          "bounds": [0, 24],
          "value_type": "int",
          "log_scale": False,
      },
      "kernel_size_a": {
          "name": "kernel_size_a",
          "type": "range",
          "bounds": [0, 24],
          "value_type": "int",
          "log_scale": False,
      },
      "kernel_size_l": {
          "name": "kernel_size_l",
          "type": "range",
          "bounds": [0, 24],
          "value_type": "int",
          "log_scale": False,
      },
      "attn_dropout": {
          "name": "attn_dropout",
          "type": "range",
          "bounds": [0.0, 0.3],
          "value_type": "float",
          "log_scale": False,
      },
      "attn_dropout_a": {
          "name": "attn_dropout_a",
          "type": "range",
          "bounds": [0.0, 0.5],
          "value_type": "float",
          "log_scale": False,
      },
      "attn_dropout_v": {
          "name": "attn_dropout_v",
          "type": "range",
          "bounds": [0.0, 0.5],
          "value_type": "float",
          "log_scale": False,
      },
      "attn_dropout_l": {
          "name": "attn_dropout_l",
          "type": "range",
          "bounds": [0.0, 0.5],
          "value_type": "float",
          "log_scale": False,
      },
      "embed_dropout": {
          "name": "embed_dropout",
          "type": "range",
          "bounds": [0.0, 0.3],
          "value_type": "float",
          "log_scale": False,
      },
      "out_dropout": {
          "name": "out_dropout",
          "type": "range",
          "bounds": [0, 0.3],
          "value_type": "float",
          "log_scale": False,
      },
      "lr": {
          "name": "lr",
          "type": "range",
          "bounds": [1e-4, 5e-3],
          "value_type": "float",
          "log_scale": True,
      },
      "nlayers": {
          "name": "nlayers",
          "type": "range",
          "bounds": [3, 7],
          "value_type": "int",
          "log_scale": False,
      },
      "num_heads": {
          "name": "num_heads",
          "type": "range",
          "bounds": [2, 12],
          "value_type": "int",
          "log_scale": False,
      },
      "dims_per_head": {
          "name": "dims_per_head",
          "type": "range",
          "bounds": [2, 12],
          "value_type": "int",
          "log_scale": False,
      },
  }

  ax_hyparams_list = [
      ax_hyparams_config_base.get(hyp_param_name)
      for hyp_param_name in search_hyp_params_list
  ]
  return ax_hyparams_list


if __name__ == "__main__":
  hyp_params = {'l_len': 20, 'a_len': 400, 'v_len': 500}
  ax_hyparams(SimpleNamespace(**hyp_params), ['lr', 'attn_dropout'])
